import json, gzip, sys
from nextapp.ont_qc.config import COLNAMES


class Info(object):
    def __init__(self):
        self.length_list = list()
        self.qscore_list = list()
        self.reads_num = 0
        self.base_num = 0

    def sort(self):
        self.length_list.sort()
        self.qscore_list.sort()

    def merge(self, info):
        self.length_list += info.length_list
        self.qscore_list += info.qscore_list
        self.reads_num += info.reads_num
        self.base_num += info.base_num

    def add(self, length, qscore):
        self.length_list.append(length)
        self.qscore_list.append(qscore)
        self.reads_num += 1
        self.base_num += length

    @property
    def stat(self):
        return {
            'read_num': self.reads_num, 'base_num': self.base_num,
            'max_len': max(self.length_list), 'min_len': min(self.length_list),
            'medium_len': self.length_list[int(self.reads_num / 2)],
            'mean_len': float('%.2f' % (self.base_num / self.reads_num)),
            'mean_qscore': float('%.2f' % (sum(self.qscore_list) / self.reads_num))
        }

    @property
    def nx0(self):
        levels = range(1, 10)
        length_dict, read_dict, base_dict = dict(), dict(), dict()
        reads_tmp, base_tmp = 0, 0
        for length in self.length_list:
            base_tmp += length
            reads_tmp += 1
            level = int(10 * base_tmp / self.base_num)
            key = 'n%d0' % level
            if level in levels and length_dict.get(key) is None:
                length_dict[key] = length
                read_dict[key] = [reads_tmp, float('%.4f' % (reads_tmp / self.reads_num))]
                base_dict[key] = [base_tmp, float('%.4f' % (base_tmp / self.base_num))]
        return {'length': length_dict, 'read': read_dict, 'base': base_dict}

    def interval(self, step, endpoint):
        data = {i: {'reads_count': 0, 'bases_sum': 0} for i in range(0, endpoint, step)}
        data[endpoint] = {'reads_count': 0, 'bases_sum': 0}
        for length in self.length_list:
            if length % step == 0:
                level = (int(length / step) - 1) * step
            else:
                level = int(length / step) * step
            if data.get(level) is None:
                data[endpoint]['reads_count'] += 1
                data[endpoint]['bases_sum'] += length
            else:
                data[level]['reads_count'] += 1
                data[level]['bases_sum'] += length
        #sys.stdout.write('base num:%d, reads num: %d\n' % (self.base_num, self.reads_num))
        out = list()
        for k, v in sorted(data.items(), key=lambda x: -x[0]):
            rc_rate, bs_rate = (v['reads_count'] / self.reads_num, v['bases_sum'] / self.base_num)
            name = ('%dK-%dK' % (k / 1000, (k + step) / 1000)) if k < endpoint else ('>%dK' % (endpoint / 1000))
            out.append({'name': name, 'reads_count': v['reads_count'], 'rc_rate': rc_rate, 'bases_sum': v['bases_sum'], 'bs_rate': bs_rate})
        return out

    def cumulative(self, step, endpoint):
        data = {i: {'reads_count': 0, 'bases_sum': 0} for i in range(0, endpoint, step)}
        data[endpoint] = {'reads_count': 0, 'bases_sum': 0}
        for length in self.length_list:
            for l in data.keys():
                if length >= l:
                    data[l]['reads_count'] += 1
                    data[l]['bases_sum'] += length
        return [{'name': ('>=%dK' % (k / 1000)),
                 'reads_count': v['reads_count'],
                 'rc_rate': v['reads_count'] / self.reads_num,
                 'bases_sum': v['bases_sum'],
                 'bs_rate': v['bases_sum'] / self.base_num} for k, v in sorted(data.items(), key=lambda x: -x[0])]


def read_summary4stat(summary4stat_file, qscore_cutoff):
    '''
    :param summary_file: run_filter 产生 summary2stat文件
    :return: [全部reads长度list, HQ高质量reads长度list, LQ低质量reads长度list]
    '''
    info_all, info_hq, info_lq = Info(), Info(), Info()
    fp = gzip.open(summary4stat_file, 'rb')
    for line in fp:
        line = line.strip().decode()
        # if not line or line.startswith(COLNAMES[0]): continue
        if not line or line.startswith('read_id'): continue
        read_id, seq_len, qscore = line.split('\t')
        seq_len, qscore = int(seq_len), float(qscore)
        info_all.add(length=seq_len, qscore=qscore)
        if qscore >= qscore_cutoff:
            info_hq.add(length=seq_len, qscore=qscore)
        else:
            info_lq.add(length=seq_len, qscore=qscore)
    fp.close()
    info_all.sort()
    info_hq.sort()
    info_lq.sort()
    return info_all, info_hq, info_lq
